import asyncio
import random
import akshare as ak
import numpy as np
from fastapi import Depends
from sqlalchemy.ext.asyncio import AsyncSession
from sqlmodel import SQLModel, Session, select

from db.session import get_collection, AsyncSessionLocal
from model.stocks import *
import pandas as pd
import datetime

from model.trades import TradeHistory, OrderHistory


async def get_stockInfoList():
    """
    获取股票信息列表，若数据库为空则从 akshare 拉取并写入数据库。
    :return: (list[dict], pandas.DataFrame)
    """
    # 获取文档
    collection = get_collection("stock_info")

    # 查询数据库
    stock_qs = await collection.find({}).sort("code").to_list(length=10000)
    if not stock_qs:
        print('数据库中不存在股票信息列表，更新数据...')
        await asyncio.sleep(random.uniform(1, 3))
        df = ak.stock_info_a_code_name()
        stock_dicts = []

        collection.insert_many([StockInfo(code=row['code'],
                                                    name=row['name']).dict()
                                          for _, row in df.iterrows()])

        for _, row in df.iterrows():
            stock = StockInfo(code=row['code'], name=row['name'])
            stock_dicts.append(stock.dict())
    else:
        stock_dicts = [StockInfo(code=s['code'], name=s['name']).dict() for s in stock_qs]

    # 构建 DataFrame
    stock_df = pd.DataFrame(stock_dicts)
    if not stock_dicts:
        return None, None
    return stock_dicts, stock_df


async def get_stockIndexList(index_symbol='000300'):
    """
    获得成分股信息列表
    :param index_symbol:
    :return: list
    """
    # 获取文档
    collection = get_collection("stock_index")

    # 查询数据库
    stock_qs = await collection.find({}).to_list(length=2000)


    if not stock_qs:
        print(f'数据库中不存在成分股信息-{index_symbol}，更新数据...')
        await asyncio.sleep(random.uniform(1, 3))
        # 从ak获取
        df = ak.index_stock_cons(index_symbol)

        stock_dicts = []

        # 插入到数据库中
        await collection.insert_many([StockIndex(
                index = index_symbol,
                code = row['品种代码'],
                name = row['品种名称'],
                date = pd.to_datetime(row['纳入日期'])
            ).dict()
            for _, row in df.iterrows()]
        )

        stock_dicts.extend([
            StockIndex(
                index=index_symbol,
                code=row['品种代码'],
                name=row['品种名称'],
                date=pd.to_datetime(row['纳入日期']))
            for _, row in df.iterrows()
        ])
    else:
        stock_dicts = [StockIndex(
                index = index_symbol,
                code = s['code'],
                name = s['name'],
                date = pd.to_datetime(s['date'])
            )
        for s in stock_qs]

    # 构建 DataFrame
    stock_df = pd.DataFrame(stock_dicts)
    if not stock_dicts:
        return None, None
    return stock_dicts, stock_df

async def get_stockPriceList(code: str,
                             start_date:str =None,
                             end_date:str =None,
                             window: int=None):
    """
    获得单个股票的行情数据 使用时间序列数据库
    :param code: 股票代码
    :param start_date: 开始日期 默认为20210101
    :param end_date: 结束日期
    :param window: 时间窗口(日) 最近的n日
    :return:list(StockPrice) pd.dataframe
    """

    async with AsyncSessionLocal() as session:

        # 加载交易日
        calendar = ak.tool_trade_date_hist_sina()
        calendar['trade_date'] = pd.to_datetime(calendar['trade_date'])

        if window is None:
            if start_date is None:
                start_date = ak.stock_individual_info_em(symbol=code)['value'][7]
            if end_date is None:
                end_date = datetime.datetime.today().date().strftime('%Y%m%d')
        else:
            # 获得近window交易日的数据
            today = datetime.datetime.today().date() # 今天的日期
            trade_days = calendar[calendar['trade_date'] < pd.to_datetime(today)] # 小于今天的交易日
            if window == 0:
                start_date = trade_days['trade_date'].iloc[-1].strftime('%Y%m%d')  # 昨天
            else:
                start_date = trade_days['trade_date'].iloc[-window].strftime('%Y%m%d') # 前n天
            end_date = trade_days['trade_date'].iloc[-1].strftime('%Y%m%d') # 昨天


        start_date_df = pd.to_datetime(start_date)
        end_date_df = pd.to_datetime(end_date)

        # 过滤交易日
        work_date = calendar[(calendar['trade_date'] >= start_date_df) & (calendar['trade_date'] <= end_date_df)]
        work_date_sets = set(work_date['trade_date'])

        # SQL 查询：根据股票代码和日期范围查找数据
        query = select(StockPrice).where(
            StockPrice.code == code,
            StockPrice.date >= start_date_df,
            StockPrice.date <= end_date_df
        )
        qs = await session.execute(query)
        stock_qs = qs.scalars().all()
        stock_dicts = []
        # 可能差几天，或者没有，找出差的日期
        # 为空的情况
        if not stock_qs:
            missing_dates_df = [start_date_df, end_date_df]
        else:
            db_date_set = set(pd.to_datetime([s.date for s in stock_qs]))
            missing_dates_df = sorted(work_date_sets - db_date_set)
            stock_dicts = [s.dict() for s in stock_qs]

        if missing_dates_df:
            start_date = missing_dates_df[0].strftime('%Y%m%d')
            end_date = missing_dates_df[-1].strftime('%Y%m%d')
            start_date_df = pd.to_datetime(start_date)
            end_date_df = pd.to_datetime(end_date)
            # 删除重复的信息
            query = select(StockPrice).where(
        StockPrice.code == code,
                StockPrice.date >= start_date_df,
                StockPrice.date <= end_date_df
            )
            qs = await session.execute(query)
            stock_qs = qs.scalars().all()
            for stock in stock_qs:
                await session.delete(stock)  # 删除每个映射对象
            await session.commit()
            await asyncio.sleep(random.uniform(1,3))
            # 从ak中获取
            df = ak.stock_zh_a_hist(
                symbol=code,
                period='daily',
                start_date=start_date,
                end_date=end_date,
            )
            if not df.empty:
                # 批量插入数据
                for _, row in df.iterrows():
                    stock_price = StockPrice(
                        date=pd.to_datetime(row['日期']),
                        code=code,
                        open=row['开盘'],
                        close=row['收盘'],
                        high=row['最高'],
                        low=row['最低'],
                        volume=row['成交量'],
                        amount=row['成交额'],
                        amplitude=row['振幅'],
                        price_change_percentage=row['涨跌幅'],
                        price_change=row['涨跌额'],
                        turnover_rate=row['换手率'],
                        weekday=pd.to_datetime(row['日期']).weekday()
                    )
                    session.add(stock_price)
                    stock_dicts.append(stock_price.dict())
                await session.commit()
                # 更新 stock_dicts
                print("获取数据成功!!!")

        # 去掉没用的字段
        stock_dicts = [{k: v for k, v in d.items() if k not in {'id', 'time'}} for d in stock_dicts]
        # 根据日期排序
        stock_dicts = sorted(stock_dicts, key=lambda x: x['date'])
        if not stock_dicts:
            return None, None
        # 构建 DataFrame
        stock_df = pd.DataFrame(stock_dicts)
        stock_df.set_index('date', inplace=True)
        return stock_dicts, stock_df

async def get_stockChipList(code: str, window: int):
    """
    获取股票的筹码数据 使用时间序列数据库
    :param code: 股票代码
    :param window: 时间窗口
    :return:
    """

    async with AsyncSessionLocal() as session:
        # 加载交易日
        calendar = ak.tool_trade_date_hist_sina()
        calendar['trade_date'] = pd.to_datetime(calendar['trade_date'])

        # 获得近window交易日的数据
        today = datetime.datetime.today().date()  # 今天的日期
        trade_days = calendar[calendar['trade_date'] < pd.to_datetime(today)]  # 小于今天的交易日
        start_date = trade_days['trade_date'].iloc[-window].strftime('%Y%m%d')  # 前n天
        end_date = trade_days['trade_date'].iloc[-1].strftime('%Y%m%d')  # 昨天

        start_date_df = pd.to_datetime(start_date)
        end_date_df = pd.to_datetime(end_date)
        # 过滤交易日
        work_date = calendar[(calendar['trade_date'] >= start_date_df) & (calendar['trade_date'] <= end_date_df)]
        work_date_sets = set(work_date['trade_date'])

        # 从获得筹码
        query = select(StockChip).where(
            StockChip.code == code,
            StockChip.date >= start_date_df,
            StockChip.date <= end_date_df
        )
        qs = await session.execute(query)
        stock_qs = qs.scalars().all()
        # 用来存储数据
        chip_dicts = []
        if not stock_qs:
            missing_dates_df = [start_date_df, end_date_df]
        else:
            db_date_set = set(pd.to_datetime([s.date for s in stock_qs]))
            missing_dates_df = sorted(work_date_sets - db_date_set)
            chip_dicts = [s.dict() for s in stock_qs]

        if missing_dates_df:
            start_date = missing_dates_df[0].strftime('%Y%m%d')
            end_date = missing_dates_df[-1].strftime('%Y%m%d')
            start_date_df = pd.to_datetime(start_date)
            end_date_df = pd.to_datetime(end_date)
            print(f'数据库中不存在该股票筹码信息-{code}: {start_date}---{end_date}，更新数据...')
            # 删除重复的信息
            query = select(StockChip).where(
                StockChip.code == code,
                StockChip.date >= start_date_df,
                StockChip.date <= end_date_df
            )
            qs = await session.execute(query)
            stock_qs = qs.scalars().all()
            for stock in stock_qs:
                await session.delete(stock)  # 删除每个映射对象
            await session.commit()
            await asyncio.sleep(random.uniform(1, 3))
            # 从ak获得筹码数据
            df = ak.stock_cyq_em(
                symbol=code,
                adjust=''
            )
            if not df.empty:
                # 获得对应的收盘价
                _, stock_df = await get_stockPriceList(
                    code = code,
                    start_date=start_date,
                    end_date=end_date,
                )
                if not stock_df is None:
                    # 取收盘价
                    close_df = stock_df.close.copy()

                    # 取最新window天
                    df = df[(pd.to_datetime(df['日期']) >= start_date_df) & (pd.to_datetime(df['日期']) <= end_date_df)].copy()

                    for index, row in df.iterrows():
                        date = pd.to_datetime(row['日期'])
                        matched = close_df.loc[date]
                        # 如果 matched 已经是 float，直接赋值
                        if isinstance(matched, (float, np.floating)):
                            close = matched
                        # 如果是 Series 或 DataFrame
                        elif isinstance(matched, pd.Series):
                            if 'close' in matched:
                                close = matched['close']
                            else:
                                close = matched.values[0]
                        elif isinstance(matched, pd.DataFrame):
                            close = matched['close'].values[0]
                        else:
                            print(f"未知类型：{type(matched)}，跳过")
                            continue


                        chip = StockChip(
                            date=date,
                            code=code,
                            close=close,
                            profit_rate=row['获利比例'],
                            avg=row['平均成本'],
                            cost_90_L=row['90成本-低'],
                            cost_90_H=row['90成本-高'],
                            concentration_90=row['90集中度'],
                            cost_70_L=row['70成本-低'],
                            cost_70_H=row['70成本-高'],
                            concentration_70=row['70集中度']
                        )
                        session.add(chip)
                        chip_dicts.append(chip.dict())
                    await session.commit()
                    print("获取数据成功!!!")

        # 去掉没用的字段
        chip_dicts = [{k: v for k, v in d.items() if k not in {'id', 'time'}} for d in chip_dicts]
        # 根据日期排序
        chip_dicts = sorted(chip_dicts, key=lambda x: x['date'])
        if not chip_dicts:
            return None, None
        # 构建 DataFrame
        chip_df = pd.DataFrame(chip_dicts)
        chip_df.set_index('date', inplace=True)
        return chip_dicts, chip_df


async def get_stockHistList(code: str, start_date: str = None, end_date: str = None, window: int = None,period='15'):
    """
    获取分时数据列表 从开盘时间 9: 30 算起 到 end_date的收盘
    :param code: 股票代码
    :param start_date: 开始日期
    :param end_date: 结束日期
    :param window: 时间窗口
    :param period: 周期 choice of {'1', '5', '15', '30', '60'}
    :return: hist_dicts(list), hist_df(dataframe) (时间     开盘     收盘     最高     最低    成交量    成交额     均价)
    """
    # 处理时间
    if window is not None:
        # 加载交易日
        calendar = ak.tool_trade_date_hist_sina()
        calendar['trade_date'] = pd.to_datetime(calendar['trade_date'])

        # 获得近window交易日的数据
        today = datetime.datetime.today().date()  # 今天的日期
        trade_days = calendar[calendar['trade_date'] < pd.to_datetime(today)]  # 小于今天的交易日
        start_date = trade_days['trade_date'].iloc[-window]  # 前n天
        end_date = trade_days['trade_date'].iloc[-1]  # 昨天
    else:
        start_date = datetime.datetime.strptime(start_date, "%Y%m%d")
        end_date = datetime.datetime.strptime(end_date, "%Y%m%d")
    start_date = start_date.strftime('%Y-%m-%d') + " 09:00:00"
    end_date = end_date.strftime("%Y-%m-%d") + " 15:00:00"
    # 从ak中查询信息
    df = ak.stock_zh_a_hist_min_em(symbol=code,
                                         start_date=start_date,
                                         end_date=end_date,
                                         period=period)
    # 用来存储数据
    hist_dicts = []

    if not df.empty:
        for index, row in df.iterrows():
            date = pd.to_datetime(row['时间'])

            hist = StockHist(
                date_time=date,
                code=code,
                date=date.date(),
                time=date.time(),
                news=row['收盘'],
                open=row['开盘'],
                close=row['收盘'],
                high=row['最高'],
                low=row['最低'],
                price_change_percentage=row['涨跌幅'],
                price_change=row['涨跌额'],
                volume=row['成交量'],
                amount=row['成交额'],
                amplitude=row['振幅'],
                turnover_rate=row['换手率']
            )

            hist_dicts.append(hist.dict())
        # 根据日期排序
        hist_dicts = sorted(hist_dicts, key=lambda x: x['date'])
    if not hist_dicts:
        return None, None
    hist_df = pd.DataFrame(hist_dicts)
    hist_df.set_index('date_time', inplace=True)
    return hist_dicts, hist_df


async def get_tradeDateList(window, date=None):
    """
    获得近window的交易日 如果date不为None, 取的则是date的时间窗口,否则为今天的时间窗口
    :param window: 时间窗口
    :param date: 日期 datetime.date类型
    :return: 返回当天+过去window天的交易日期 index为-1的就是当天的日期 ['20250703', '20250704', '20250707', '20250708', '20250709', '20250710', '20250711', '20250714', '20250715', '20250716', '20250717', '20250718', '20250721', '20250722', '20250723']
    """
    # 获取交易日历（返回的是 datetime.date 类型）
    df = ak.tool_trade_date_hist_sina()

    if date is not None:
        today = date
    else:
        # 过滤掉今天之后的日期
        today = datetime.datetime.today().date()

    df = df[df['trade_date'] <=today]

    # 取最近 window 天，按时间倒序，然后再升序返回
    df = df.sort_values('trade_date', ascending=False).head(window+1)
    trade_days = [d.strftime('%Y%m%d') for d in df['trade_date']]
    return sorted(trade_days)


async def get_stockBidList(code):
    """
    获得股票实时行情
    :param code: 股票代码
    :return: dict, series
    """

    df = ak.stock_bid_ask_em(symbol=code)
    if not df.empty:

        df = df.set_index('item').copy().T
        # 尝试将所有列转换为 float，如果不能转换的则为 NaN
        df = df.apply(pd.to_numeric, errors='coerce')
        df = df.fillna(0)
        bid = StockBid(code=code,
                       sell_5=df['sell_5'].values[0],
                       sell_5_vol=df['sell_5_vol'].values[0],
                       sell_4=df['sell_4'].values[0],
                       sell_4_vol=df['sell_4_vol'].values[0],
                       sell_3=df['sell_3'].values[0],
                       sell_3_vol=df['sell_3_vol'].values[0],
                       sell_2=df['sell_2'].values[0],
                       sell_2_vol=df['sell_2_vol'].values[0],
                       sell_1=df['sell_1'].values[0],
                       sell_1_vol=df['sell_1_vol'].values[0],
                       buy_1=df['buy_1'].values[0],
                       buy_1_vol=df['buy_1_vol'].values[0],
                       buy_2=df['buy_2'].values[0],
                       buy_2_vol=df['buy_2_vol'].values[0],
                       buy_3=df['buy_3'].values[0],
                       buy_3_vol=df['buy_3_vol'].values[0],
                       buy_4=df['buy_4'].values[0],
                       buy_4_vol=df['buy_4_vol'].values[0],
                       buy_5=df['buy_5'].values[0],
                       buy_5_vol=df['buy_5_vol'].values[0],
                       news=df['最新'].values[0],
                       average_price=df['均价'].values[0],
                       raise_rate=df['涨幅'].values[0],
                       raise_fall=df['涨跌'].values[0],
                       total_price=df['总手'].values[0],
                       amount=df['金额'].values[0],
                       turnover=df['换手'].values[0],
                       quantity_ratio=df['量比'].values[0],
                       high=df['最高'].values[0],
                       low=df['最低'].values[0],
                       today_open=df['今开'].values[0],
                       yesterday_close=df['昨收'].values[0],
                       limit_up=df['涨停'].values[0],
                       limit_down=df['跌停'].values[0],
                       external=df['外盘'].values[0],
                       inner=df['内盘'].values[0]
                       )
        series = pd.Series(bid.dict(), name=bid.code)
        return bid.dict(), series

    return None, None


async def insert_tradeHistory(tradeHistory):
    # 存入数据库中
    collection = get_collection("trade_history")
    collection.insert_one(tradeHistory.dict())

async def get_tradeHistory(orderID):
    # 获取文档
    collection = get_collection("trade_history")
    # 查询tradeHistory
    trade_qs = await collection.find({"orderID": orderID}).sort("date_time", 1).to_list(length=1000)
    if trade_qs == []:
        return None
    trade_dicts = [TradeHistory(orderID=t['orderID'],
                                date_time=t['date_time'],
                                price=t['price'],
                                signal=t['signal'],
                                profit_cash=t['profit_cash'],
                                cash=t['cash'],
                                new_position=t['new_position'],
                                total_position=t['total_position']) for t in trade_qs]
    return trade_dicts

async def insert_orderHistory(orderHistory):
    # 获得文档
    collection = get_collection("order_history")
    # 插入orderHistory
    collection.insert_one(orderHistory.dict())

async def get_orderHistory():
    """
    获取订单记录
    :return: [{orderID , date_time}]
    """
    collection = get_collection("order_history")
    order_qs = await collection.find({}).to_list()
    if order_qs == []:
        return None
    order_dicts = [OrderHistory(orderID=o['orderID'],
                                code=o['code'],
                                date_time=o['date_time'],
                                strategy_name=o['strategy_name'],
                                evaluate=o['evaluate']) for o in order_qs]
    return order_dicts

async def delete_history(orderID):
    """
    删除订单记录
    :param orderID: 订单ID
    :return:
    """
    # 获得文档
    collection = get_collection("trade_history")
    # 获得原始文档数量
    init_count = await collection.count_documents({})
    # 删除交易记录
    await collection.delete_many({"orderID": orderID})
    if await collection.count_documents({}) == init_count:
        return None
    # 获取文档
    collection = get_collection("order_history")
    # 获得原始文档数量
    init_count = await collection.count_documents({})
    # 删除订单记录
    await collection.delete_one({"orderID": orderID})
    if await collection.count_documents({}) == init_count:
        return None

    return True

async def delete_allHistory():

    collection = get_collection("trade_history")
    collection.delete_many({})

    collection = get_collection("order_history")
    collection.delete_many({})

